from sklearn.datasets import (
    load_iris, load_wine, load_breast_cancer, load_digits
)
from sklearn import tree
from sklearn.model_selection import train_test_split

# 加载UCI数据集，划分训练集和测试集
iris_X, iris_y = load_iris(return_X_y=True)
iris_X_train, iris_X_test, iris_y_train, iris_y_test = train_test_split(
    iris_X, iris_y, test_size=0.3, random_state=0
)

# 生成未剪枝和剪枝后两个决策树
clf_iris = tree.DecisionTreeClassifier().fit(iris_X_train, iris_y_train)
clf_iris_pruned = tree.DecisionTreeClassifier(max_depth=3).fit(iris_X_train, iris_y_train)

print('\niris_dataset :')
print('剪枝前错误率：', 1 - clf_iris.score(iris_X_test, iris_y_test))
print('剪枝后错误率：', 1 - clf_iris_pruned.score(iris_X_test, iris_y_test))


wine_X, wine_y = load_wine(return_X_y=True)
wine_X_train, wine_X_test, wine_y_train, wine_y_test = train_test_split(
    wine_X, wine_y, test_size=0.3, random_state=0
)

clf_wine = tree.DecisionTreeClassifier().fit(wine_X_train, wine_y_train)
clf_wine_pruned = tree.DecisionTreeClassifier(max_depth=4).fit(wine_X_train, wine_y_train)

print('\nwine_dataset :')
print('剪枝前错误率：', 1 - clf_wine.score(wine_X_test, wine_y_test))
print('剪枝后错误率：', 1 - clf_wine_pruned.score(wine_X_test, wine_y_test))


breast_cancer_X, breast_cancer_y = load_breast_cancer(return_X_y=True)
breast_cancer_X_train, breast_cancer_X_test, breast_cancer_y_train, breast_cancer_y_test = train_test_split(
    breast_cancer_X, breast_cancer_y, test_size=0.3, random_state=0
)

clf_breast_cancer = tree.DecisionTreeClassifier().fit(breast_cancer_X_train, breast_cancer_y_train)
clf_breast_cancer_pruned = tree.DecisionTreeClassifier(max_depth=6).fit(breast_cancer_X_train, breast_cancer_y_train)

print('\nbreast_cancer_dataset :')
print('剪枝前错误率：', 1 - clf_breast_cancer.score(breast_cancer_X_test, breast_cancer_y_test))
print('剪枝后错误率：', 1 - clf_breast_cancer_pruned.score(breast_cancer_X_test, breast_cancer_y_test))


digits_X, digits_y = load_digits(return_X_y=True)
digits_X_train, digits_X_test, digits_y_train, digits_y_test = train_test_split(
    digits_X, digits_y, test_size=0.3, random_state=0
)

clf_digits = tree.DecisionTreeClassifier().fit(digits_X_train, digits_y_train)
clf_digits_pruned = tree.DecisionTreeClassifier(max_depth=12).fit(digits_X_train, digits_y_train)

print('\ndigits_dataset :')
print('剪枝前错误率：', 1 - clf_digits.score(digits_X_test, digits_y_test))
print('剪枝后错误率：', 1 - clf_digits_pruned.score(digits_X_test, digits_y_test))
